from detectron2.config import get_cfg
from detectron2.modeling import GeneralizedRCNN
from detectron2.model_zoo import model_zoo


class FastRCNNOdetector(GeneralizedRCNN):
    def __init__(self):
        labels_zh = ['jdk', 'jzp', 'wsz', 'xjp', 'nnh']
        # labels_zh = ['jzp']
        CONFIG_PATH = "COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml"
        cfg = get_cfg()
        cfg.merge_from_file(model_zoo.get_config_file(CONFIG_PATH))
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
        cfg.MODEL.ROI_HEADS.NUM_CLASSES = len(labels_zh)
        cfg.MODEL.DEVICE = "cpu"
        # cfg.MODEL.WEIGHTS = "./output/model_final.pth"
        super(FastRCNNOdetector, self).__init__(cfg=cfg)
